#ifndef __PKT_DROP__
#define __PKT_DROP__

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/ip.h>
#include <linux/icmp.h>
#include <linux/udp.h>
#include <linux/in.h>
#include <linux/moduleparam.h>
#include <linux/string.h>
#include <linux/netdevice.h>
#include <linux/fs.h>
#include <linux/version.h>
#include <asm/uaccess.h>

// ------------------------------------------------------------------
// Kernel module definitions
#define AUTHOR 		"Stefano Paris"
#define DESC   		"Packet Dropping"
#define DESC_LONG  	"This module performs the packet dropping"

static struct nf_hook_ops netfilter_ops_pre;

char BACKBONE_INTERFACES[256] = "eth0";

#define DEF_DROP_DEGREE  0
#define DEF_DROP_WND    10

// Default destination port used by ping
#define IP_PING_PROT  1

static uint32_t rx_uni_pck_count  = 0;
static uint32_t drop_wnd_index    = 0;
static uint32_t drop_count        = 0;
static uint32_t drop_deg          = DEF_DROP_WND;

static uint32_t drop_wnd[DEF_DROP_WND];

MODULE_AUTHOR(AUTHOR);
MODULE_DESCRIPTION(DESC " - " DESC_LONG);
MODULE_LICENSE("GPL");
module_param(drop_deg, int, 0600);
MODULE_PARM_DESC(drop_deg, "Ratio of dropped pakets (x/10 dropped, x = drop_deg)");

// ------------------------------------------------------------------
/*
 * Pre-Routing function: drop the packet sent by other nodes
 */
unsigned int hook_pre_routing(unsigned int hooknum,
			       struct sk_buff *skb,
			       const struct net_device *in,
			       const struct net_device *out,
			       int (*okfn)(struct sk_buff *)) {
	
	u_int8_t icmp_type = 0;
	uint8_t *ip_payload = NULL;
	struct iphdr *iph = (struct iphdr *)skb_network_header(skb);
	
	if (!drop_deg) {
		// No dropping
		return NF_ACCEPT;
	} 
	//printk(">>>>>> %u\n", iph->protocol);
	if (iph->protocol == IPPROTO_ICMP) {
		ip_payload = (skb->data + (iph->ihl * 4));
		memcpy(&icmp_type, ip_payload, 1);
		//printk(">>>>>> %u\n", icmp_type);
		if ((icmp_type == ICMP_ECHO) &&
			(drop_wnd[drop_wnd_index] == 1)) {
			// Drop the rx packet
			//printk(">>>>>> DROP %u\n", icmp_type);
			drop_wnd_index = ++rx_uni_pck_count % DEF_DROP_WND;
			drop_count++;
			return NF_DROP;
		}
	}
	drop_wnd_index = ++rx_uni_pck_count % DEF_DROP_WND;
	
	return NF_ACCEPT; /* Accept ALL Packets */
}

// ------------------------------------------------------------------
/*
 * Initialization function
 */
int init_test(void) {
	
	int ret = 0;
	int i;
		
	printk(KERN_ALERT "Loading module %s...\n", DESC);
	
	// Initialize dropping filter
	if (drop_deg > DEF_DROP_WND) {
		printk(KERN_ALERT "Sorry, dropping degree (%d) cannot be grater than DEF_DROP_WND (%d)\n", drop_deg, DEF_DROP_WND);
		return -1;
	}
	memset(drop_wnd, 0, DEF_DROP_WND * sizeof(uint32_t));
	for (i = 0; i < drop_deg; i++) {
		drop_wnd[i] = 1;
	}
	
	printk(KERN_ALERT "   BackBone Interfaces: %s\n", BACKBONE_INTERFACES);
	
	// Register the filtering function
	netfilter_ops_pre.hook     =  hook_pre_routing;
	netfilter_ops_pre.pf       =  PF_INET;
	netfilter_ops_pre.hooknum  =  NF_INET_PRE_ROUTING;
	netfilter_ops_pre.priority =  NF_IP_PRI_FIRST;
	
	ret = nf_register_hook(&netfilter_ops_pre); /* register NF_IP_PRE_ROUTING hook */
	if (ret < 0) {
		printk(KERN_ALERT "Sorry, registering PRE_ROUTING hook registered failed with %d \n", ret);
		return ret;
	}
	printk(KERN_ALERT "   PRE_ROUTING hook registered (%d)!\n", ret);	
	
	// Initialize and show the internal state variables
	rx_uni_pck_count  = 0;
	drop_count   = 0;
	
	printk(KERN_ALERT "   Dropping degree  = %d\n", drop_deg);
	printk(KERN_ALERT "   Tot Rx packets   = %d\n", rx_uni_pck_count);
	printk(KERN_ALERT "   Dropped packets  = %d\n", drop_count);
	
	return 0;
}

/*
 * Exit function
 */
void cleanup_test(void) {
	
	printk(KERN_ALERT "Unloading module %s...\n", DESC);
	
	// Unregister the filtering function
	nf_unregister_hook(&netfilter_ops_pre);
	memset(&netfilter_ops_pre, 0, sizeof(struct nf_hook_ops));
	//kfree(netfilter_ops_pre);
	printk(KERN_ALERT "   PRE_ROUTING hook released.\n");
		
	// Show the internal state variables
	printk(KERN_ALERT "   Dropping degree  = %d\n", drop_deg);
	printk(KERN_ALERT "   Tot Rx packets   = %d\n", rx_uni_pck_count);
	printk(KERN_ALERT "   Dropped packets  = %d\n", drop_count);
	
	// Erase dropping filter
	memset(drop_wnd, 0, DEF_DROP_WND * sizeof(uint32_t));
	
}

module_init(init_test);
module_exit(cleanup_test);

#endif
